/*
 <one line to give the program's name and a brief idea of what it does.>
 Copyright (C) <year>  <name of author>
 
 This program is free software: you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published by
 the Free Software Foundation, either version 3 of the License, or
 (at your option) any later version.
 
 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU General Public License for more details.
 
 You should have received a copy of the GNU General Public License
 along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 */

#include "msLagrangian.h"


namespace atomism
{
    
    
    bool msLagrangian::isLagrangianRegisteredInPython=0;
    msRegistrar msLagrangian::Registrar("msLagrangian", msLagrangian::createInstance);
    
    ostream& operator<<(ostream& out,const msLagrangian& motion) {  motion.print(out); return out; }
    
    //-------------------------------------------------------------------------------------------------
    //-------------------------------------------------------------------------------------------------
    
    
    void msLagrangian::registryInPython()
    {
#if USE_PYTHON
        msMotion::registryInPython();
        
        if( ! msLagrangian::isLagrangianRegisteredInPython )
        {
            using namespace boost::python;
            
            class_<msLagrangian,bases<msMotion>,boost::shared_ptr<msLagrangian> >
            ("msLagrangian",
             "Defines a mechanical system. ",no_init )
            .def( "New", &msLagrangian::New , "Create a new object")
            .staticmethod("New")
            .def("setEqOfMotion" ,  &msLagrangian::setEqOfMotion , "defines the system. Arg1: kinetic operator; Arg2: Potential energy" )
            .def("E" ,  (double(msLagrangian::*)() ) &msLagrangian::E , "Total energy at the current state" )
            .def("T" ,  (double(msLagrangian::*)() ) &msLagrangian::T , "Kinetic energy at the current state" )
            .def("U" ,  (double(msLagrangian::*)() ) &msLagrangian::U , "Potential energy at the current state" )
            .def("L" ,  (double(msLagrangian::*)() ) &msLagrangian::L , "Action at the current state" )
            .def(self_ns::str(self_ns::self));
            msLagrangian::isLagrangianRegisteredInPython=1;
        }
#endif
    }
    
    //-------------------------------------------------------------------------------------------------
    //-------------------------------------------------------------------------------------------------
    
    msLagrangian::msQIntegrand::msQIntegrand(const msLagrangian* sys,double temp) : msScalarFunction()
    {
        system = sys;   T = temp;
        R     = system->getUnits()->convert( msUnit::J_mol , csts::R  );
        Econv = system->getUnits()->convert( system->Epot->getYUnit(), 1);
        
        fact   = pow( 2*csts::pi*double(csts::kb)*T , double(system->GeneralizedCoordinates->getIntegrationDim()/2.) )
        / pow( double(csts::h) ,system->GeneralizedCoordinates->getIntegrationDim() );
    }
    
    //-------------------------------------------------------------------------------------------------
    
    double msLagrangian::msQIntegrand::evaluate()
    {
        double V   =  system->Epot->evaluate() * Econv;
        double Kin =  system->KineticOperator->kineticFunction( system->SiUnits );
        double r   =  fact * exp( - V  /( R * T )) * sqrt( Kin );
        return r;
    }
    
    //-------------------------------------------------------------------------------------------------
    //-------------------------------------------------------------------------------------------------
    
    msLagrangian::msQbyObservable::msQbyObservable(msObservable* observable, const msLagrangian* sys,double temp)
    : msQIntegrand(sys,temp) { Observable = observable; sumQ=0; }
    
    double msLagrangian::msQbyObservable::evaluate()
    {
        double q   = msLagrangian::msQIntegrand::evaluate();
        sumQ 	   += q;
        double obs = (*Observable)();
        return q * obs;
    }
    
    //-------------------------------------------------------------------------------------------------
    //-------------------------------------------------------------------------------------------------
    
    msLagrangian::msDOSIntegrand::msDOSIntegrand(const msLagrangian* sys,double ener)
    {
        system  = sys;
        Econv   = system->getUnits()->convert( system->Epot->getYUnit() , 1 );
        EconvSi = system->SiUnits.convert( system->getUnits()->Energy , 1 );
        
        ndof = system->GeneralizedCoordinates->getIntegrationDim( );
        E    = ener;
        emin = system->getSampler()->sample(*(system->Epot)).Min;
        fact = pow( 2*csts::pi / pow( double(csts::h) ,2 ), double(ndof/2.) )
        / Gamma(ndof/2.);
    }
    
    //-------------------------------------------------------------------------------------------------
    
    double msLagrangian::msDOSIntegrand::evaluate()
    {
        double Kin =  system->KineticOperator->kineticFunction( system->SiUnits );
        double V   =  system->Epot->evaluate() * Econv ; // in System->UserUnits
        
        if(V<emin) LOGGER_WRITE_ALONE(msLogger::WARNING,"e < emin","double msLagrangian::msDOSIntegrand::Evaluate() ");
        
        double DOS=0;
        if( E - ( V - emin) > E/1000.)
        DOS = fact * sqrt(Kin) * pow( ( E - ( V - emin) ) * EconvSi, ndof/2.-1);
        
        return DOS;
    }
    
    //-------------------------------------------------------------------------------------------------
    // Class msLagrangian
    //-------------------------------------------------------------------------------------------------
    
    boost::shared_ptr<msTreeMapper>  msLagrangian::setEqOfMotion( boost::shared_ptr<msKineticOperator> kin ,
                                                                 boost::shared_ptr<msScalarFunction> epot ) {
        
        LOGGER_ENTER_FUNCTION_DBG("void msLagrangian::defineEqOfMotion(...)",getFullId());
        msTreeMapper::update(KineticOperator,kin);
        msTreeMapper::update(Epot,epot);
        
        if(    Epot->getCoordinates().get()
           != KineticOperator->getCoordinates().get() ) {
            
            BOOST_THROW_EXCEPTION( msError("The coordinates of the PES and kinetic operators have to point on the same variables"
                                           ,"boost::shared_ptr<msTreeMapper>  msLagrangian::setEqOfMotion( boost::shared_ptr<msKineticOperator> kin , boost::shared_ptr<msScalarFunction> epot )"
                                           ,getFullId() ));
        }
        if(!kin->hasParent())  kin->setAffiliation(mySharedPtr());
        if(!epot->hasParent()) epot->setAffiliation(mySharedPtr());
        
        msTreeMapper::update(GeneralizedCoordinates,KineticOperator->getCoordinates());
        
        if(!GeneralizedCoordinates->hasParent()) GeneralizedCoordinates->setAffiliation(mySharedPtr());
        
        update();
        LOGGER_EXIT_FUNCTION();
        return mySharedPtr();
    }
    
    //-------------------------------------------------------------------------------------------------
    //-------------------------------------------------------------------------------------------------
    
    double msLagrangian::DOS(double E) {
        
        LOGGER_ENTER_FUNCTION("double msLagrangian::DOS(double T) const",getFullId());
        LOGGER_WRITE(msLogger::INFO,"Sample the generalized coordinates space");
        
        msDOSIntegrand DOSfct(this,E);
        DOSfct.initialize();
        DOSfct.setYUnit("J^-1");
        DOSfct.setCoordinates(GeneralizedCoordinates.getSharedPtr());
        double DOS = 0;
        
        try{ DOS =  Sampler->sample( DOSfct , SiUnits ).Integral;
        }
        catch( msError& e ) {
            e.addContext("configurational space integration for DOS calculation failed (double msLagrangian::DOS(double E) ");
            throw;
        }
        
        DOS *= 1/getUnits()->convert(msUnit::J,1);
        LOGGER_EXIT_FUNCTION();
        return DOS * Mult / Sym;
    }
    
    //-------------------------------------------------------------------------------------------------
    //-------------------------------------------------------------------------------------------------
    
    double msLagrangian::Q(double T) {
        
        LOGGER_ENTER_FUNCTION("double msLagrangian::Q(double T) const",getFullId());
        LOGGER_WRITE(msLogger::INFO,"Sample the generalized coordinates space");
        
        double Q = 0;
        
        msQIntegrand Qfct(this,T);
        Qfct.initialize();
        Qfct.setYUnit("");
        Qfct.setCoordinates(GeneralizedCoordinates.getSharedPtr());
        
        try{ Q =  Sampler->sample( Qfct , SiUnits ).Integral;
        }
        catch( msError& e ) {
            e.addContext("configurational space integration for partition function calculation failed (double msLagrangian::Q(double T)");
            throw;
        }
        LOGGER_EXIT_FUNCTION();
        return Q * Mult / Sym;
    }
    
    //-------------------------------------------------------------------------------------------------
    //-------------------------------------------------------------------------------------------------
    
    double msLagrangian::averageByBoltzmann( msObservable& observable , double T )  {
        
        LOGGER_ENTER_FUNCTION("double msLagrangian::averageByBoltzmann( msObservable& observable , double T ) ",getFullId());
        LOGGER_WRITE(msLogger::INFO,"Sample the generalized coordinates space");
        
        double average = 0;
        
        msQbyObservable observe( &observable , this, T );
        observe.initialize();
        observe.setYUnit("");
        observe.setCoordinates(GeneralizedCoordinates.getSharedPtr());
        
        try{
            average = Sampler->sample(observe).Integral / ( observe.sumQ * GeneralizedCoordinates->getVolumeDq() ) ;
        }
        catch( msError& e ) {
            
            e.addContext("configurational space integration for observable calculation failed (double msLagrangian::averageByBoltzmann( msObservable& observable , double T )) ");
            throw;
        }
        LOGGER_EXIT_FUNCTION();
        return average * Mult / Sym;
    }
    
    //-------------------------------------------------------------------------------------------------
    //-------------------------------------------------------------------------------------------------
    /*
     double msLagrangian::getOriginEnergy() const
     {
     return OriginEnergy;
     }
     double msLagrangian::scanMinOfPes()
     {
     return getUnits()->convert( Epot->getYUnit() , Sampler->sample(*Epot).Min );
     }
     
     */
    
    //-------------------------------------------------------------------------------------------------
    //-------------------------------------------------------------------------------------------------
    
    double msLagrangian::E() { GeneralizedCoordinates->getUnfreezedValues(SiUnits,q);
        return getUnits()->convert(SiUnits.Energy ,E( q , qp));
    }
	
    double msLagrangian::U() { GeneralizedCoordinates->getUnfreezedValues(SiUnits,q);
        return getUnits()->convert(SiUnits.Energy , U( q ));
    }
	
    double msLagrangian::T() { GeneralizedCoordinates->getUnfreezedValues(SiUnits,q);
        return getUnits()->convert(SiUnits.Energy,T( q , qp));
    }
    
    double msLagrangian::L() { GeneralizedCoordinates->getUnfreezedValues(SiUnits,q);
        return getUnits()->convert(SiUnits.Energy,L( q , qp));
    }
    
    //-------------------------------------------------------------------------------------------------
    //-------------------------------------------------------------------------------------------------
    
    double msLagrangian::T(vector_type q_, vector_type qp_) {  GeneralizedCoordinates->setUnfreezedValues(SiUnits,q=q_);
        KineticOperator->computeKMat( SiUnits );
        vector_type tmp = KineticOperator->getKMatrix()*qp_;
        return 0.5*tmp.dot(qp_);
    }
    //-------------------------------------------------------------------------------------------------
    double msLagrangian::U(vector_type q_) {  GeneralizedCoordinates->setUnfreezedValues(SiUnits,q_);
        return  SiUnits.convert( Epot->getYUnit(), Epot->evaluate());
    }
    //-------------------------------------------------------------------------------------------------
    double msLagrangian::L(vector_type q_, vector_type qp_) { return T(q_, qp_)-U(q_);  //J/mol
    }
    
    double msLagrangian::L(vector_type q_, vector_type qp_,bool compkmat) {
        
        if(compkmat) return L(q_,qp_);
        
        vector_type tmp = KineticOperator->getKMatrix()*qp_;
        double L = 0.5*tmp.dot(qp_) + U(q_);
        return L;
    }
    
    //-------------------------------------------------------------------------------------------------
    double msLagrangian::E(vector_type q_, vector_type qp_)  { return T(q_, qp_) + U(q_);    //J/mol
    }
    //-------------------------------------------------------------------------------------------------
    void msLagrangian::computeAll(vector_type q_, vector_type qp_,double& E,double& Ep,double& Ek,double& L) {
        
        Ek=T(q_, qp_);
        Ep=U(q_);
        E=Ek+Ep;
        L=Ek-Ep;
    }
    
    //-------------------------------------------------------------------------------------------------
    //-------------------------------------------------------------------------------------------------
    
    std::ostream& msLagrangian::print(std::ostream& out) const {
        
        msMotion::print(out);
        output::printHeader(out,"Lagrangian");
        
        out<<"Potential energy: "<< Epot->getFullId() <<endl;
        out<<"Kinetic operator: "<< KineticOperator->getFullId() <<endl;
        out<<"\nId of the generalized coordinates:  "<<GeneralizedCoordinates->getId()<<"; type: "<<GeneralizedCoordinates->getType()<<endl<<"\t";
        GeneralizedCoordinates->getUnits()->abstract(out);
        out<<"\n\tList of unfreezed coordinates:\n";
        msGeneralizedCoordinates::UnfreezIt qi = GeneralizedCoordinates->beginUnfreezed();
        
        for(; qi!=GeneralizedCoordinates->endUnfreezed(); qi++) {
            out<<"\t\t"; (*qi).abstract(out); out<<endl;}
        
        // out<<"\nCurrent state in SI:\nq:\n"<<q<<"\nqp: \n"<<qp<<endl;
        return out;
    }
    
    
    
    /*
     testok
     if( GeneralizedCoordinates.getSharedPtr() == boost::shared_ptr<msGeneralizedCoordinates>() )  
     BOOST_THROW_EXCEPTION( msError("GeneralizedCoordinates member not initialized (did you call msLagrangian::initialize()?)"
     ,"double msLagrangian::Q(double T)",getFullId()) );
     
     if( Sampler.getSharedPtr() == boost::shared_ptr<msSampler>() )  
     BOOST_THROW_EXCEPTION( msError("Sampler member not initialized (did you call msLagrangian::initialize()?)"
     ,"double msLagrangian::Q(double T)",getFullId()) );
     
     */
    
    
}
